"""
215. 数组中的第K个最大元素
中等
相关标签
premium lock icon
相关企业
给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。



示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5
示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4


提示：

1 <= k <= nums.length <= 105
-104 <= nums[i] <= 104
"""
import random
from typing import List


class Solution:
    def findKthLargest(self, nums, k):
        """
        快排类似
        :param nums:
        :param k:
        :return:
        """
        def find(nums,k):
            if len(nums)==1:
                return nums[0]
            ch = random.choice(nums)
            small = []
            equals = []
            large = []
            for num in nums:
                if ch>num:
                    small.append(num)
                elif ch<num:
                    large.append(num)
                else:
                    equals.append(num)
            if k>len(small) and k <= len(small) + len(equals):
                return ch
            elif k<=len(small):
                return find(small,k)
            else:
                return find(large,k - len(small) - len(equals))
        return find(nums,len(nums) - k + 1)

    def findKthLargest_res1(self, nums, k):
        def quick_select(nums, k):
            # 随机选择基准数
            pivot = random.choice(nums)
            big, equal, small = [], [], []
            # 将大于、小于、等于 pivot 的元素划分至 big, small, equal 中
            for num in nums:
                if num > pivot:
                    big.append(num)
                elif num < pivot:
                    small.append(num)
                else:
                    equal.append(num)
            if k <= len(big):
                # 第 k 大元素在 big 中，递归划分
                return quick_select(big, k)
            if len(nums) - len(small) < k:
                # 第 k 大元素在 small 中，递归划分
                return quick_select(small, k - len(nums) + len(small))
            # 第 k 大元素在 equal 中，直接返回 pivot
            return pivot

        return quick_select(nums, k)


# 测试用例
if __name__ == "__main__":
    solution = Solution()

    # 测试用例 1
    result1 = solution.findKthLargest([3, 2, 1, 5, 6, 4], 2)
    print(f"测试用例 1: {result1}")  # 预期输出: 5

    # 测试用例 2
    result2 = solution.findKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)
    print(f"测试用例 2: {result2}")  # 预期输出: 4

    # 测试用例 3
    result3 = solution.findKthLargest([1], 1)
    print(f"测试用例 3: {result3}")  # 预期输出: 1

    # 测试用例 4
    result4 = solution.findKthLargest([7, 10, 4, 3, 20, 15], 3)
    print(f"测试用例 4: {result4}")  # 预期输出: 10

    # 测试用例 5
    result5 = solution.findKthLargest([-1, -1], 2)
    print(f"测试用例 5: {result5}")  # 预期输出: -1

